#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版键盘检测模块 - 更精确的键盘检测算法
"""

import cv2
import numpy as np
import numpy.typing as npt
from typing import Tuple, List, Dict, Optional, Any
import logging
from dataclasses import dataclass
from enum import Enum
import math

class KeyIssueType(Enum):
    MISSING = "缺失"
    MISPLACED = "位置错误"
    ROTATED = "方向错误"
    WRONG_KEY = "错误按键"
    SIZE_MISMATCH = "大小不匹配"

@dataclass
class KeyIssue:
    issue_type: KeyIssueType
    position: Tuple[int, int]
    description: str
    confidence: float

class EnhancedKeyboardDetector:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
        # 优化的检测参数
        self.min_keyboard_area = 30000  # 降低最小键盘区域面积
        self.key_contour_area_range = (50, 8000)  # 扩大按键轮廓面积范围
        self.aspect_ratio_tolerance = 0.4  # 增加长宽比容差
        
        # 更严格的匹配参数
        self.match_threshold = 0.6  # 降低匹配阈值，提高敏感度
        self.position_tolerance = 30  # 增加位置容差
        self.angle_tolerance = 15  # 增加角度容差
        self.size_tolerance = 0.3  # 大小容差
        
        # 多尺度检测参数
        self.scale_factors = [0.8, 1.0, 1.2]  # 多尺度检测
        self.rotation_angles = [0, 90, 180, 270]  # 检测旋转角度
        
        # 特征匹配参数
        self.orb = cv2.ORB_create(nfeatures=1000)
        self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        
    def detect_keyboard_keys(self, image: npt.NDArray) -> List[Dict[str, Any]]:
        """
        增强版按键检测 - 使用多种方法提高检测精度
        """
        if image is None:
            return []
            
        # 预处理图像
        processed_images = self._multi_scale_preprocess(image)
        
        all_keys = []
        for scale, processed_img in processed_images.items():
            # 方法1: 轮廓检测
            contour_keys = self._detect_by_contours(processed_img, scale)
            
            # 方法2: 模板匹配
            template_keys = self._detect_by_template_matching(processed_img, scale)
            
            # 方法3: 特征点检测
            feature_keys = self._detect_by_features(processed_img, scale)
            
            # 合并结果
            all_keys.extend(contour_keys)
            all_keys.extend(template_keys)
            all_keys.extend(feature_keys)
        
        # 去重和优化
        final_keys = self._merge_and_optimize_keys(all_keys, image)
        
        # 按位置排序
        final_keys.sort(key=lambda k: (k['y'], k['x']))
        
        return final_keys
    
    def _multi_scale_preprocess(self, image: npt.NDArray) -> Dict[float, npt.NDArray]:
        """多尺度预处理"""
        processed_images = {}
        
        for scale in self.scale_factors:
            # 缩放图像
            if scale != 1.0:
                h, w = image.shape[:2]
                new_h, new_w = int(h * scale), int(w * scale)
                scaled_image = cv2.resize(image, (new_w, new_h))
            else:
                scaled_image = image.copy()
            
            # 预处理
            processed = self._preprocess_for_key_detection(scaled_image)
            processed_images[scale] = processed
            
        return processed_images
    
    def _preprocess_for_key_detection(self, image: npt.NDArray) -> npt.NDArray:
        """增强版图像预处理"""
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
        
        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)
        
        # 双边滤波去噪
        filtered = cv2.bilateralFilter(enhanced, 9, 75, 75)
        
        # 多种阈值方法
        # 方法1: 自适应阈值
        adaptive_thresh = cv2.adaptiveThreshold(
            filtered, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, 11, 2
        )
        
        # 方法2: Otsu阈值
        _, otsu_thresh = cv2.threshold(filtered, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        
        # 方法3: Canny边缘检测
        edges = cv2.Canny(filtered, 50, 150)
        
        # 合并多种方法的结果
        combined = cv2.bitwise_or(adaptive_thresh, otsu_thresh)
        combined = cv2.bitwise_or(combined, edges)
        
        # 形态学操作
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        cleaned = cv2.morphologyEx(combined, cv2.MORPH_CLOSE, kernel)
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel)
        
        return cleaned
    
    def _detect_by_contours(self, binary_image: npt.NDArray, scale: float) -> List[Dict[str, Any]]:
        """基于轮廓的按键检测"""
        contours, _ = cv2.findContours(
            binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        
        keys = []
        for contour in contours:
            area = cv2.contourArea(contour)
            scaled_area_range = (
                self.key_contour_area_range[0] * scale * scale,
                self.key_contour_area_range[1] * scale * scale
            )
            
            if scaled_area_range[0] <= area <= scaled_area_range[1]:
                # 检查轮廓特征
                if self._is_valid_key_contour(contour):
                    key_info = self._extract_key_info(contour, scale)
                    if key_info:
                        keys.append(key_info)
        
        return keys
    
    def _detect_by_template_matching(self, image: npt.NDArray, scale: float) -> List[Dict[str, Any]]:
        """基于模板匹配的按键检测"""
        # 这里可以实现模板匹配逻辑
        # 暂时返回空列表，可以后续扩展
        return []
    
    def _detect_by_features(self, image: npt.NDArray, scale: float) -> List[Dict[str, Any]]:
        """基于特征点的按键检测"""
        # 检测特征点
        keypoints, descriptors = self.orb.detectAndCompute(image, None)
        
        keys = []
        for kp in keypoints:
            # 将特征点转换为按键信息
            key_info = {
                'x': int(kp.pt[0] / scale),
                'y': int(kp.pt[1] / scale),
                'width': int(kp.size / scale),
                'height': int(kp.size / scale),
                'angle': kp.angle,
                'area': int((kp.size / scale) ** 2),
                'confidence': kp.response,
                'scale': scale,
                'detection_method': 'feature'
            }
            keys.append(key_info)
        
        return keys
    
    def _is_valid_key_contour(self, contour: npt.NDArray) -> bool:
        """检查轮廓是否为有效的按键轮廓"""
        # 计算轮廓特征
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)
        
        if perimeter == 0:
            return False
        
        # 计算圆度
        circularity = 4 * np.pi * area / (perimeter * perimeter)
        
        # 计算长宽比
        x, y, w, h = cv2.boundingRect(contour)
        aspect_ratio = w / h if h > 0 else 0
        
        # 计算凸性
        hull = cv2.convexHull(contour)
        hull_area = cv2.contourArea(hull)
        convexity = area / hull_area if hull_area > 0 else 0
        
        # 检查条件
        return (0.3 <= circularity <= 1.0 and  # 圆度合理
                0.3 <= aspect_ratio <= 3.0 and  # 长宽比合理
                convexity >= 0.7)  # 凸性足够
    
    def _extract_key_info(self, contour: npt.NDArray, scale: float) -> Optional[Dict[str, Any]]:
        """提取按键信息"""
        # 获取边界框
        x, y, w, h = cv2.boundingRect(contour)
        
        # 计算最小外接矩形
        rect = cv2.minAreaRect(contour)
        center, size, angle = rect
        
        # 计算轮廓特征
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)
        
        key_info = {
            'x': int(center[0] / scale),
            'y': int(center[1] / scale),
            'width': int(size[0] / scale),
            'height': int(size[1] / scale),
            'angle': angle,
            'area': area / (scale * scale),
            'perimeter': perimeter / scale,
            'confidence': min(1.0, area / 1000),  # 基于面积的置信度
            'scale': scale,
            'detection_method': 'contour',
            'contour': contour
        }
        
        return key_info
    
    def _merge_and_optimize_keys(self, all_keys: List[Dict], original_image: npt.NDArray) -> List[Dict[str, Any]]:
        """合并和优化检测到的按键"""
        if not all_keys:
            return []
        
        # 按置信度排序
        all_keys.sort(key=lambda k: k.get('confidence', 0), reverse=True)
        
        final_keys = []
        used_positions = []
        
        for key in all_keys:
            # 检查是否与已有按键重叠
            is_duplicate = False
            for used_pos in used_positions:
                distance = np.sqrt(
                    (key['x'] - used_pos[0])**2 + (key['y'] - used_pos[1])**2
                )
                if distance < 20:  # 重叠阈值
                    is_duplicate = True
                    break
            
            if not is_duplicate:
                final_keys.append(key)
                used_positions.append((key['x'], key['y']))
        
        return final_keys
    
    def compare_keyboards(self, reference_keys: List[Dict], current_keys: List[Dict]) -> List[KeyIssue]:
        """增强版键盘对比 - 修复检测逻辑"""
        issues = []
        
        # 为每个参考按键找到最佳匹配
        used_current_indices = set()
        matched_reference_keys = []
        
        for ref_idx, ref_key in enumerate(reference_keys):
            best_match_idx = None
            best_match_score = 0
            best_match_distance = float('inf')
            
            for i, curr_key in enumerate(current_keys):
                if i in used_current_indices:
                    continue
                
                # 计算位置距离
                distance = np.sqrt(
                    (ref_key['x'] - curr_key['x'])**2 + 
                    (ref_key['y'] - curr_key['y'])**2
                )
                
                # 计算相似度
                similarity = self._calculate_enhanced_similarity(ref_key, curr_key)
                
                # 综合评分 - 优先考虑位置匹配
                if distance < self.position_tolerance:
                    if distance < best_match_distance:
                        best_match_idx = i
                        best_match_score = similarity
                        best_match_distance = distance
            
            if best_match_idx is not None:
                used_current_indices.add(best_match_idx)
                curr_key = current_keys[best_match_idx]
                matched_reference_keys.append((ref_idx, best_match_idx))
                
                # 检查具体问题
                issues.extend(self._check_enhanced_key_issues(ref_key, curr_key))
            else:
                # 未找到匹配的按键 - 这是缺失按键
                issues.append(KeyIssue(
                    issue_type=KeyIssueType.MISSING,
                    position=(ref_key['x'], ref_key['y']),
                    description=f"按键缺失 - 位置 ({ref_key['x']}, {ref_key['y']})",
                    confidence=0.9
                ))
        
        # 检查是否有额外的按键（不应该存在的按键）
        for i, curr_key in enumerate(current_keys):
            if i not in used_current_indices:
                issues.append(KeyIssue(
                    issue_type=KeyIssueType.WRONG_KEY,
                    position=(curr_key['x'], curr_key['y']),
                    description=f"多余按键 - 位置 ({curr_key['x']}, {curr_key['y']})",
                    confidence=0.8
                ))
        
        return issues
    
    def _calculate_enhanced_similarity(self, key1: Dict, key2: Dict) -> float:
        """计算增强版按键相似度"""
        # 位置相似度
        pos_distance = np.sqrt(
            (key1['x'] - key2['x'])**2 + (key1['y'] - key2['y'])**2
        )
        pos_similarity = max(0, 1 - pos_distance / 100)
        
        # 大小相似度
        size_diff = abs(key1['width'] - key2['width']) + abs(key1['height'] - key2['height'])
        size_similarity = max(0, 1 - size_diff / 100)
        
        # 角度相似度
        angle_diff = abs(key1['angle'] - key2['angle']) % 360
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
        angle_similarity = max(0, 1 - angle_diff / 90)
        
        # 面积相似度
        area_diff = abs(key1['area'] - key2['area']) / max(key1['area'], key2['area'])
        area_similarity = max(0, 1 - area_diff)
        
        # 综合相似度（加权平均）
        total_similarity = (
            pos_similarity * 0.4 + 
            size_similarity * 0.25 + 
            angle_similarity * 0.2 + 
            area_similarity * 0.15
        )
        
        return total_similarity
    
    def _check_enhanced_key_issues(self, ref_key: Dict, curr_key: Dict) -> List[KeyIssue]:
        """检查增强版按键问题"""
        issues = []
        
        # 检查位置偏差
        pos_distance = np.sqrt(
            (ref_key['x'] - curr_key['x'])**2 + (ref_key['y'] - curr_key['y'])**2
        )
        
        if pos_distance > self.position_tolerance:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.MISPLACED,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键位置偏差 {pos_distance:.1f} 像素",
                confidence=min(0.9, pos_distance / 50)
            ))
        
        # 检查旋转
        angle_diff = abs(ref_key['angle'] - curr_key['angle']) % 360
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
            
        if angle_diff > self.angle_tolerance:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.ROTATED,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键旋转 {angle_diff:.1f} 度",
                confidence=min(0.9, angle_diff / 45)
            ))
        
        # 检查大小差异
        width_diff = abs(ref_key['width'] - curr_key['width']) / ref_key['width']
        height_diff = abs(ref_key['height'] - curr_key['height']) / ref_key['height']
        
        if width_diff > self.size_tolerance or height_diff > self.size_tolerance:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.SIZE_MISMATCH,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键大小异常 (宽差: {width_diff:.1%}, 高差: {height_diff:.1%})",
                confidence=max(width_diff, height_diff)
            ))
        
        return issues
    
    def visualize_detection(self, image: npt.NDArray, keys: List[Dict], issues: List[KeyIssue], reference_keys: List[Dict] = None) -> npt.NDArray:
        """可视化检测结果 - 改进版"""
        vis_image = image.copy()
        
        # 绘制检测到的按键（绿色框）
        for key in keys:
            # 绘制按键边界框
            cv2.rectangle(
                vis_image,
                (key['x'] - key['width']//2, key['y'] - key['height']//2),
                (key['x'] + key['width']//2, key['y'] + key['height']//2),
                (0, 255, 0), 2  # 绿色 - 检测到的按键
            )
            
            # 绘制按键中心点
            cv2.circle(vis_image, (key['x'], key['y']), 3, (0, 0, 255), -1)
        
        # 绘制问题标记
        missing_count = 0
        other_issues_count = 0
        
        for issue in issues:
            x, y = issue.position
            
            if issue.issue_type == KeyIssueType.MISSING:
                # 缺失按键 - 绘制红色框
                color = (0, 0, 255)  # 红色
                symbol = "MISSING"
                missing_count += 1
                
                # 绘制红色矩形框
                cv2.rectangle(
                    vis_image,
                    (x - 25, y - 25),
                    (x + 25, y + 25),
                    color, 3
                )
                
                # 绘制对角线
                cv2.line(vis_image, (x - 25, y - 25), (x + 25, y + 25), color, 2)
                cv2.line(vis_image, (x - 25, y + 25), (x + 25, y - 25), color, 2)
                
            elif issue.issue_type == KeyIssueType.MISPLACED:
                color = (0, 255, 255)  # 黄色 - 位置错误
                symbol = "MISPLACED"
                other_issues_count += 1
                
            elif issue.issue_type == KeyIssueType.ROTATED:
                color = (255, 255, 0)  # 青色 - 旋转
                symbol = "ROTATED"
                other_issues_count += 1
                
            elif issue.issue_type == KeyIssueType.SIZE_MISMATCH:
                color = (255, 0, 255)  # 紫色 - 大小不匹配
                symbol = "SIZE_MISMATCH"
                other_issues_count += 1
                
            else:
                color = (128, 0, 128)  # 深紫色 - 其他错误
                symbol = "OTHER"
                other_issues_count += 1
            
            # 绘制问题标记文字
            cv2.putText(
                vis_image, symbol, (x - 30, y - 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2
            )
            
            # 绘制问题圆圈（非缺失按键）
            if issue.issue_type != KeyIssueType.MISSING:
                cv2.circle(vis_image, (x, y), 15, color, 2)
        
        # 添加统计信息
        total_keys = len(keys)
        total_issues = len(issues)
        
        # 绘制信息面板
        cv2.rectangle(vis_image, (10, 10), (400, 120), (0, 0, 0), -1)
        cv2.putText(vis_image, f"Detected Keys: {total_keys}", (20, 35),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        cv2.putText(vis_image, f"Missing Keys: {missing_count}", (20, 60),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        cv2.putText(vis_image, f"Other Issues: {other_issues_count}", (20, 85),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
        cv2.putText(vis_image, f"Total Issues: {total_issues}", (20, 110),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        
        return vis_image
